import os
import pynvml
import warnings
#! 设置显卡
pynvml.nvmlInit() #! 初始化
deviceCount = pynvml.nvmlDeviceGetCount() #! 统计显卡数量
GPUs = []
for busid in range(deviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(busid)
    device = pynvml.nvmlDeviceGetMemoryInfo(handle)
    memory = device.total // 1048576 / 1024 #! 单位: G
    if memory > 8: GPUs.append(str(busid))
pynvml.nvmlShutdown() #! 关闭
# GPUs = ['0'] #! 显存多的卡
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(GPUs) #! 设置显卡设备
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" #! 设置警告ERROR级别
#! 导入官方包
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
#! 导入自定义模块
from model import SSD300, SSD512
from loss import SSDSSDLoss, YOLOSSDLoss
from data import PascalVOCDatasetV1, PascalVOCDatasetV2
#! 忽略警告类报错
warnings.filterwarnings('ignore')

#! 数据集设置
data_folder = 'dataset/VOCdevkit'
voc_labels = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
              'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor')
label_map = {k: v + 1 for v, k in enumerate(voc_labels)}
label_map['background'] = 0
n_classes = len(label_map)
rev_label_map = {v:k for k,v in label_map.items()}
PascalVOCDataset = PascalVOCDatasetV2
output_shape = (300,300)
batch_size   = 16
train_loader = PascalVOCDataset(batch_size, data_folder, 'TRAIN', output_shape=output_shape, shuffle=True)
valid_loader = PascalVOCDataset(batch_size, data_folder, 'TEST', output_shape=output_shape)

#! 训练设置
SSD  = SSD300
Loss = SSDSSDLoss
lr = 1e-2 #! 学习率 Adam 建议 lr = 1e-3
decay_lr_to = 0.065 #! 0.065/0.07 0.1:7 0.065:10
momentum = 0.9  # momentum
weight_decay = 1e-4  # weight decay
epochs = 500
grad_clip = True
model_path = None
retraining = True         #! 重新训练
continue_training = False #! 续训练

#! 开始训练
if model_path is None:
    SSDModel = SSD(n_classes=n_classes)
    lr_schedule = keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=lr,  #! 初始学习率
        decay_rate=1-weight_decay,  #! 衰减率
        decay_steps=len(train_loader),
    )
    optimizer = keras.optimizers.SGD(learning_rate=lr_schedule, momentum=momentum)
    SSDModel.model.compile(optimizer=optimizer)
else:
    SSDModel = SSD(n_classes=n_classes,model_path=model_path)
    if retraining and not continue_training:
        lr_schedule = keras.optimizers.schedules.ExponentialDecay(
            initial_learning_rate=lr,  #! 初始学习率
            decay_rate=1-weight_decay,  #! 衰减率
            decay_steps=len(train_loader),
        )    
        optimizer = keras.optimizers.SGD(learning_rate=lr_schedule, momentum=momentum)
        SSDModel.model.compile(optimizer=optimizer)
        SSDModel.layer_freeze(14) #! 取消全部冻结, 所有层均可训练
    elif continue_training and not retraining:
        pass
    else:
        raise ValueError("retraining or continue_training only one can be set to True")

#! 损失函数
criterion = Loss(SSDModel.priors_cxcy)

SSDModel.train(train_loader, SSDModel.model.optimizer, criterion, epochs, valid_loader=train_loader, eval_epochs=100, save_epochs=50)